import numpy as np

from task import TASK


class ROTATION(TASK):
    
    
    def __init__(self) -> None:
        super().__init__()
        
    
    def get_fitness(self, position, orientation):
        
        if (position[0] ** 2 + position[1] ** 2) ** 0.5 < 0.5:
            self.__fitness = orientation[2]
        else:
            self.__fitness = -((position[0] ** 2 + position[1] ** 2) ** 0.5)
        
        return self.__fitness
    
        
    def is_finished(self, position, orientation, target):
        
        target_theta = np.arctan((position[1] - target[1]) / (position[0] - target[0]))
        if np.abs(orientation[2] - target_theta) <= np.pi / 18:
            return True
        
        else:
            return False


    def save_network(self, index):
        
        if self.network is not None:
            np.save(f"networks/rotation/{index}-{self.__fitness:.3f}", self.network)
